package com.atguigu.ad.hive;

import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ConstantObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;
import org.apache.hadoop.io.IOUtils;
import org.lionsoul.ip2region.xdb.Searcher;

import java.io.IOException;
import java.util.ArrayList;

/**
 * @author Hefei
 * @description 自定义解析IP
 * @project_name com.atguigu.ad.hive
 * @since 2023/4/6 14:14
 */

@Description(name = "parse_ip",
        value = "_FUNC_(\"path\", \"ip\") - Custom IP parsing function " +
                "\n  that uses the ip2region.xdb file " +
                "\n  passed into HDFS to parse IP addresses into a structure " +
                "\n  similar to" +
                "\n  {\"country\":\"China\",\"area\":\"0\",\"province\":\"Beijing\",\"city\":\"Beijing\",\"isp\":\"China YiDong\"}" +
                "\nThe function accepts two parameters: " +
                "\n  the first is the path to the HDFS file, " +
                "\n  and the second is the IP address. " +
                "\nPassing in invalid parameters will result in an error.",
        extended = "For example: " +
                "\n > select parse_ip(\"hdfs://hadoop102:8020/ip2region/ip2region.xdb\",\"120.245.112.30\");" +
                "\nThe result is: " +
                "\n > {\"country\":\"China\",\"area\":\"0\",\"province\":\"Beijing\",\"city\":\"Beijing\",\"isp\":\"China YiDong\"}")
public class ParseIP extends GenericUDF {
    private transient PrimitiveObjectInspector.PrimitiveCategory[] inputTypes = new PrimitiveObjectInspector.PrimitiveCategory[2];
    private Searcher searcher;

    /**
     * 初始化 规定我们返回值的数据类型
     *
     * @param arguments The ObjectInspector for the arguments
     * @return
     * @throws UDFArgumentException
     */
    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {

        // todo1.我们需要对传入的参数进行校验，参数个数、参数类型

        // 检查参数列表的个数
        checkArgsSize(arguments, 2, 2);

        // 检查参数列表中传入的参数的类型
        checkArgPrimitive(arguments, 0);
        checkArgPrimitive(arguments, 1);

        // 检查参数列表中的传入的参数是否是属于String类型
        checkArgGroups(arguments, 0, inputTypes, PrimitiveObjectInspectorUtils.PrimitiveGrouping.STRING_GROUP);
        checkArgGroups(arguments, 1, inputTypes, PrimitiveObjectInspectorUtils.PrimitiveGrouping.STRING_GROUP);

        // todo 需要在此处创建出searcher对象

        if (arguments[0] instanceof ConstantObjectInspector) {
            String filePath = getConstantStringValue(arguments, 0);
            Path path = new Path(filePath);
            Configuration configuration = new Configuration();

            try {
                // 通过get方法传入conf来构建FS对象
                FileSystem fs = FileSystem.get(configuration);

                // 使用fs中的open方法传入路径 拿到输入流
                FSDataInputStream inputStream = fs.open(path);
                // 创建一个输出流
                ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
                // 使用我们hadoop中的copy方法 对两个流进行复制
                IOUtils.copyBytes(inputStream, outputStream, 1024);

                // 将输出流转成byte数组
                byte[] cBuff = outputStream.toByteArray();

                // 构建searcher对象
                searcher = Searcher.newWithBuffer(cBuff);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        // todo 2. 返回数据类型对象检查器  我们采用Struct结构来返回我们的数据
        ArrayList<String> structFieldName = new ArrayList<String>();
        ArrayList<ObjectInspector> structFieldObjectInspector = new ArrayList<ObjectInspector>();

        structFieldName.add("country");
        structFieldName.add("area");
        structFieldName.add("province");
        structFieldName.add("city");
        structFieldName.add("isp");

        structFieldObjectInspector.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        structFieldObjectInspector.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        structFieldObjectInspector.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        structFieldObjectInspector.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        structFieldObjectInspector.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);

        return ObjectInspectorFactory.getStandardStructObjectInspector(structFieldName, structFieldObjectInspector);
    }

    /**
     * 具体处理我们来的每一条数据
     *
     * @param arguments The arguments as DeferedObject, use DeferedObject.get() to get the
     *                  actual argument Object. The Objects can be inspected by the
     *                  ObjectInspectors passed in the initialize call.
     *                  参数作为 DeferedObject，使用 DeferedObject.get()获取实际参数 Object。
     *                  对象可以由初始化调用中传递的对象检查器进行检查。
     * @return
     * @throws HiveException
     */
    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {

        String ip = arguments[1].get().toString();
        String region = null;
        try {
            region = searcher.search(ip);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        String[] split = region.split("\\|");

        ArrayList<String> list = new ArrayList<String>();
        list.add(split[0]);
        list.add(split[1]);
        list.add(split[2]);
        list.add(split[3]);
        list.add(split[4]);

        return list;
    }

    /**
     * explain 来查看执行计划
     * @param children
     * @return
     */
    @Override
    public String getDisplayString(String[] children) {
        return getStandardDisplayString("parse_ip", children);
    }
}